#!/usr/bin/env python
import numpy as np

PI = np.pi

class Problem():
    def __init__(self, case=None):
        self.case = case
        self.domain = (0, 1, -1, 1)

        self.x_pde = None
        self.x_bc = None
        self.x_ic = None
        self.x_test = None

        self.target_pde = None
        self.target_bc = None
        self.target_ic = None
        self.target_test = None

        #
        self.pde_name = None
        self.eval_list_pde = None
        self.operator_type = None
        self.eq_names = None
        self.out_var = None

    def __repr__(self):
        sep = '*****************' * 3
        text_pde = f'{self.pde_name} (case={self.case}):'
        text_train = f'\tx_pde:   \t{None if self.x_pde is None else self.x_pde.shape}'
        text_test = f'\tx_test:    \t{None if self.x_test is None else self.x_test.shape}'
        text_train_bc = f'\tx_bc:    \t{None if self.x_bc is None else self.x_bc.shape}'
        text_train_ic = f'\tx_ic:    \t{None if self.x_ic is None else self.x_ic.shape}'
        return '\n'.join([sep, text_pde, text_train, text_train_bc, text_train_ic, text_test, sep])

    # def u_exact(self, x_in):
    #     # u_exact are used to check solution error and problem
    #     raise NotImplementedError('Not Implemented')

    def rhs(self, x_in):
        # right hand side of pde (forcing terms)
        raise NotImplementedError('Not Implemented')

    def bc(self, x_in):
        # right hand side of pde (forcing terms)
        raise NotImplementedError('Not Implemented')

    def lhs(self, *args, **kwargs):
        # right hand side of pde (pde operators)
        raise NotImplementedError('Not Implemented')


class Poisson_1d(Problem):
    def __init__(self, case=1):
        super().__init__(case)
        assert case in [1, 2, 3, 4, 5, 6, 7], f'case {case} not implemented'

        self.pde_name = 'Poisson'
        self.out_var = 'u'
        self.domain = (0, 2*PI)

    def u_exact(self, x_in):
        if self.case == 1:
            return np.sin(x_in)
        elif self.case == 2:
            return np.cos(2 * x_in)
        elif self.case == 3:
            return np.sin(10 * x_in)
        elif self.case == 4:
            return np.cos(10 * x_in) + np.cos(x_in)
        elif self.case == 5:
            return np.sin(100 * x_in)
        elif self.case == 6:
            return np.sin(x_in) + (x_in - 2)**2 

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        if self.case == 1:
            return - np.sin(x_in)
        elif self.case == 2:
            return - 4 * np.cos(2 * x_in)
        elif self.case == 3:
            return - 100 * np.sin(10 * x_in)
        elif self.case == 4:
            return - 100 * np.cos(10 * x_in) - np.cos(x_in)
        elif self.case == 5:
            return - 10000 * np.sin(100 * x_in)
        elif self.case == 6:
            return - np.sin(x_in) + 2
        

    def lhs(self, u_xx, u_yy):
        # Poisson operator
        return u_xx + u_yy

class Poisson_2d(Problem):
    def __init__(self, case=1):
        super().__init__(case)
        assert case in [1, 2, 3, 4], f'case {case} not implemented'

        self.pde_name = 'Poisson'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.domain = (0, 1, 0, 1)

    def u_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y)

        if self.case == 2:
            return np.cos(np.pi * x) * np.cos(np.pi * y)

        if self.case == 3:
            return x * (1 - x) * y * (1 - y) * np.exp(x - y)

        if self.case == 4:
            return np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y)

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) * (-8) * np.pi ** 2

        if self.case == 2:
            return np.cos(np.pi * x) * np.cos(np.pi * y) * (-2) * np.pi ** 2

        if self.case == 3:
            return 2 * x * (y - 1) * (y - 2 * x + x * y + 2) * np.exp(x - y)

        if self.case == 4:
            # error
            return np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) * (-32) * np.pi ** 2

    def lhs(self, u_xx, u_yy):
        # Poisson operator
        return u_xx + u_yy


class Burgers_2d(Problem):
    def __init__(self, case=1):
        super().__init__(case)
        assert case in [1, 2, 3, 4], f'case {case} not implemented'

        self.pde_name = 'Burgers'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.domain = (0, 1, -1, 1)

    def pde_rhs(self, tx):
        m = tx.shape[0]
        return np.zeros((m, 1))

    def bc_rhs(self, tx):
        m = tx.shape[0]
        return np.zeros((m, 1))

    def ic_rhs(self, tx):
        x = tx[:, [1]]
        return np.sin(-np.pi * x)